import os
import numpy as np
import pandas as pd
import joblib
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.impute import SimpleImputer
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import StandardScaler, LabelBinarizer, LabelEncoder
from sklearn.model_selection import train_test_split, cross_val_score

"""
多分类
# 线性回归模型
from sklearn.linear_model import LinearRegression

# 决策树模型
from sklearn.tree import DecisionTreeRegressor

# 随机森林
from sklearn.ensemble import RandomForestRegressor
"""


HOUSING_PATH = "datasets/kddcup"
def load_dataset(housing_path=HOUSING_PATH):
    csv_path = os.path.join(housing_path, "kddcup.data_10_percent_corrected")
    return pd.read_csv(csv_path)

class DataFrameSelector(BaseEstimator, TransformerMixin):
    def __init__(self, attribute_names):
        self.attribute_names = attribute_names

    def fit(self, X, y = None):
        return self

    def transform(self, X, y=None):
        return X[self.attribute_names].values

    def fit_transform(self, X, y=None):
        return self.transform(X)

# data的cat_attrib属性变为独热码返回，行数不变，列数变为此属性的种类数
def getOneHotCode(data, cat_attrib, whole_dataset):
    selector = DataFrameSelector(cat_attrib)
    cat_1 = selector.transform(data)
    encoder = LabelBinarizer()
    tmp = selector.transform(whole_dataset)
    encoder.fit(tmp)
    tmp = encoder.transform(cat_1)
    # print(encoder.classes_)
    return tmp

def getLabelEncode(data, cat_attrib, whole_dataset):
    selector = DataFrameSelector(cat_attrib)
    cat_1 = selector.transform(data)
    encoder = LabelEncoder()
    tmp = selector.transform(whole_dataset)
    encoder.fit(tmp.ravel())
    tmp = encoder.transform(cat_1.ravel())
    # print(encoder.classes_)
    return tmp

def full_pipeline(data, num_attribs, cat_attribs, whole_dataset):
    selector = DataFrameSelector(num_attribs)
    num_1 = selector.transform(data)
    imputer = SimpleImputer(strategy="median")
    num_2 = imputer.fit_transform(num_1)
    std_scaler = StandardScaler()
    num_3 = std_scaler.fit_transform(num_2)
    # print(num_3.shape)
    # print(num_3)

    for ele in cat_attribs:
        tmp = getOneHotCode(data, ele, whole_dataset)
        # print(type(tmp))
        # print(tmp.shape)
        # df_ = pd.DataFrame(tmp, columns=list(label_binarizer.classes_))
        # print(df_.head())
        num_3 = np.c_[num_3, tmp]
    return num_3



def model(reg, data_prepared, data_labels, test_prepared, test_labels, save_name="my_model.pkl"):
    print("\n", save_name)
    reg.fit(data_prepared, data_labels)
    test_predictions = reg.predict(test_prepared)
    print("Predictions:\t", test_predictions[:10])
    print("Labels:\t\t", list(test_labels)[:10])
    print("rmse:\t\t",mean_error(test_predictions, test_labels))

    # reg_scores = cross_val_score(reg, data_prepared, data_labels, scoring="neg_mean_squared_error",cv=10)
    # reg_rmse_scores = np.sqrt(-reg_scores)
    # display_scores(reg_rmse_scores)

    # 保存模型
    joblib.dump(reg, save_name)
    # 导入模型
    # my_model_loaded = joblib.load("my_model.pkl")


def mean_error(data_labels, data_predictions):
    mse = mean_squared_error(data_labels, data_predictions)
    rmse = np.sqrt(mse)
    return rmse

def display_scores(scores):
    print("Scores:", scores)
    print("Mean:", scores.mean())
    print("Standard deviation:", scores.std())

if __name__ == '__main__':
    kddcup = load_dataset()  # <class 'pandas.core.frame.DataFrame'>
    # print(kddcup.info())  # (494021, 42)
    labels = getLabelEncode(kddcup, ["class"], kddcup)
    kddcup["class"] = labels

    cat_attribs = ["protocol_type", "service", "flag"]  # 3
    num_attribs = list(kddcup)  # 42 - 3 - 1= 38
    num_attribs.remove("class")
    for ele in cat_attribs:
        num_attribs.remove(ele)

    # (395216, 42) (98805, 42)
    train_set, test_set = train_test_split(kddcup, test_size=0.2, random_state=42)
    train_set_labels = train_set["class"].copy()
    train_set = train_set.drop("class", axis=1)  # (395216, 41)
    test_set_labels = test_set["class"].copy()
    test_set = test_set.drop("class", axis=1)  # (98805, 41)

    # (395216, 118) (98805, 118)
    train_set_prepared = full_pipeline(train_set, num_attribs, cat_attribs, kddcup)
    test_set_prepared = full_pipeline(test_set, num_attribs, cat_attribs, kddcup)

    # 线性回归模型
    from sklearn.linear_model import LinearRegression
    lin_reg = LinearRegression()
    model(lin_reg, train_set_prepared, train_set_labels, test_set_prepared, test_set_labels, "mul_lin_model.pkl")

    # 决策树模型
    from sklearn.tree import DecisionTreeRegressor
    tree_reg = DecisionTreeRegressor()
    model(tree_reg, train_set_prepared, train_set_labels, test_set_prepared, test_set_labels, "mul_tree_model.pkl")

    # 随机森林
    from sklearn.ensemble import RandomForestRegressor
    forest_reg = RandomForestRegressor()
    model(forest_reg, train_set_prepared, train_set_labels, test_set_prepared, test_set_labels, "mul_forest_model.pkl")

